import matplotlib.pyplot as plt
from sklearn.datasets.samples_generator import make_blobs
from sklearn.cluster import MiniBatchKMeans
from sklearn import metrics
import json

class KMeans:
    def __init__(self, dataset, args):
        self.dataset = dataset


    def run(self):
        X, y = make_blobs(n_samples=1000, n_features=2, centers=[[-1,-1], [0,0], [1,1], [2,2]], cluster_std=[0.4, 0.2, 0.2, 0.2],
                  random_state =9)
        result = {'2':{}, '3':{}, '4':{}, '5':{}}

        for index, k in enumerate((2,3,4,5)):
            y_pred = MiniBatchKMeans(n_clusters=k, batch_size = 200, random_state=9).fit_predict(X)

            #score= metrics.calinski_harabaz_score(X, y_pred)
            data = self.store(k, y_pred, X)

            #result[str(k)]['score'] = round(score, 2)
            result[str(k)]['data'] = data

        return {
            'Result': result,
            'Description': '使用随机生成的二维坐标点进行聚类，分别进行了多个聚类核的拟合，并进行了结果对比及聚类分数评估。后续可以应用到食品数据中，进行食品风险的预警值聚类'
        }

    def store(self, k, y_pred, X):
        res = {}
        x = [round(_, 3) for _ in X[:, 0]]
        y = [round(_, 3) for _ in X[:, 1]]
        if k == 2:
            res['red'] = []
            res['blue'] = []
            for i in range(len(y_pred)):
                if y_pred[i] == 0:
                    res['red'].append([x[i], y[i]])
                else:
                    res['blue'].append([x[i], y[i]])
        elif k == 3:
            res['red'] = []
            res['blue'] = []
            res['green'] = []
            for i in range(len(y_pred)):
                if y_pred[i] == 0:
                    res['red'].append([x[i], y[i]])
                elif y_pred[i] == 1:
                    res['blue'].append([x[i], y[i]])
                else:
                    res['green'].append([x[i], y[i]])
        elif k == 4:
            res['red'] = []
            res['blue'] = []
            res['green'] = []
            res['yellow'] = []
            for i in range(len(y_pred)):
                if y_pred[i] == 0:
                    res['red'].append([x[i], y[i]])
                elif y_pred[i] == 1:
                    res['blue'].append([x[i], y[i]])
                elif y_pred[i] == 2:
                    res['green'].append([x[i], y[i]])
                else:
                    res['yellow'].append([x[i], y[i]])
        else:
            res['red'] = []
            res['blue'] = []
            res['green'] = []
            res['yellow'] = []
            res['pink'] = []
            for i in range(len(y_pred)):
                if y_pred[i] == 0:
                    res['red'].append([x[i], y[i]])
                elif y_pred[i] == 1:
                    res['blue'].append([x[i], y[i]])
                elif y_pred[i] == 2:
                    res['green'].append([x[i], y[i]])
                elif y_pred[i] == 3:
                    res['yellow'].append([x[i], y[i]])
                else:
                    res['pink'].append([x[i], y[i]])
        return res




if __name__ == '__main__':
    print(KMeans('../../data/k-means_dataset1', {}).run())